package com.moran.controller;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.api.Advisor;
import org.springframework.ai.document.Document;
import org.springframework.ai.rag.advisor.RetrievalAugmentationAdvisor;
import org.springframework.ai.rag.generation.augmentation.ContextualQueryAugmenter;
import org.springframework.ai.rag.retrieval.search.VectorStoreDocumentRetriever;
import org.springframework.ai.reader.TextReader;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.io.FileSystemResource;
import org.springframework.core.io.Resource;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import java.util.List;
import java.util.stream.Collectors;

@RestController
public class RagController {

    @Autowired
    private SimpleVectorStore vectorStore;

    // 原始文档路径
    private static final String DOCUMENT_FILE_PATCH = "C:\\Users\\ThinkPad\\Desktop\\document.txt";

    // 向量文档路径
    private static final String VECTOR_FILE_PATCH = "C:\\Users\\ThinkPad\\Desktop\\vector_store.json";


    private final ChatClient chatClient;

    public RagController(ChatClient.Builder chatClientBuilder) {
        this.chatClient = chatClientBuilder.build();
    }

    /**
     * 基于RAG的知识问答
     *
     * @param input
     * @return
     */
    @RequestMapping("/rag")
    public String rag(String input) {
        Advisor retrievalAugmentationAdvisor = RetrievalAugmentationAdvisor.builder()
                .documentRetriever(VectorStoreDocumentRetriever.builder()
                        .similarityThreshold(0.50)
                        .vectorStore(vectorStore)
                        .build())
                .queryAugmenter(ContextualQueryAugmenter.builder()
                         .allowEmptyContext(true) // 检索到的上下文为空时，模型用自己的方式回答用户查询，默认为false，不回答
                        .build())
                .build();

        String answer = chatClient.prompt()
                .advisors(retrievalAugmentationAdvisor)
                .user(input)
                .call()
                .content();
        return answer;
    }

    /**
     * 添加文档到向量库
     *
     * @return
     */
    @RequestMapping("/addDocument")
    public String addDocument() {
        try {
//            File vectorStoreFile = new File(VECTOR_FILE_PATCH);
//            if (vectorStoreFile.exists()) {
//                // 如果向量文档存在，直接进行加载
//                vectorStore.load(vectorStoreFile);
//            } else {
                // 读取文档
                Resource resource = new FileSystemResource(DOCUMENT_FILE_PATCH);
                TextReader textReader = new TextReader(resource);
                List<Document> documents = textReader.get();

                // 文本分割（可选，提高检索效果）
                TokenTextSplitter splitter = new TokenTextSplitter();
                List<Document> splitDocuments = splitter.apply(documents);

                // 创建向量存储并添加文档
                vectorStore.add(splitDocuments);

                // 可选：保存到本地（避免每次重启重新生成）
//                vectorStore.save(new File(VECTOR_FILE_PATCH));
//            }
        } catch (Exception e) {
            throw new RuntimeException("Failed to create vector store from document", e);
        } finally {
            return "添加成功";
        }
    }

    /**
     * 从向量库中删除文档
     *
     * @return
     */
    @RequestMapping("/delDocument")
    public String delDocument() {
        String source = DOCUMENT_FILE_PATCH.substring(DOCUMENT_FILE_PATCH.lastIndexOf("\\") + 1);
        List<Document> existingDocs = vectorStore.similaritySearch(
                SearchRequest.builder()
                        .query("null")
                        .filterExpression("source == '" + source + "'")
                        .topK(1000)
                        .build() // 设置一个足够大的值来获取所有相关文档
        );

        // 提取文档ID并删除
        List<String> idsToDelete = existingDocs.stream()
                .map(Document::getId)
                .collect(Collectors.toList());

        if (!idsToDelete.isEmpty()) {
            vectorStore.delete(idsToDelete);
        }

        return "删除成功";
    }

}
